import numpy as np
import torch
x_train = np.load("/home/wangwei83/Desktop/pytorch-dl-cv/dataset/mnist/x_train.npy")
y_train_label = np.load("/home/wangwei83/Desktop/pytorch-dl-cv/dataset/mnist/y_train_label.npy")
x = torch.tensor(y_train_label[:5],dtype=torch.int64)
# 定义一个张量输入，因为此时有 5 个数值，且最大值为9，且 类别数为 10
# 所以我们可以得到 y 的输出结果的形状应该为 shape=(5,10);5行12列
y = torch.nn.functional.one_hot(x, 10)  # 一个参数张量x, 10为类别数


import numpy as np
import torch
x_train = np.load("/home/wangwei83/Desktop/pytorch-dl-cv/dataset/mnist/x_train.npy")
y_train_label = np.load("/home/wangwei83/Desktop/pytorch-dl-cv/dataset/mnist/y_train_label.npy")
x = torch.tensor(y_train_label[:5],dtype=torch.int64)
# 定义一个张量输入，因为此时有 5 个数值，且最大值为9，且 类别数为 10
# 所以我们可以得到 y 的输出结果的形状应该为 shape=(5,10);5行12列
y = torch.nn.functional.one_hot(x, 10)  # 一个参数张量x, 10为类别数

import torch.nn as nn
class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28,312),
            nn.ReLU(),
            nn.Linear(312, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )
    def forward(self, input):
        x = self.flatten(input)
        logits = self.linear_relu_stack(x)
        return logits







